package com.tianhaochen.springai.config;

import com.tianhaochen.springai.mongo.ChatMessage;
import com.tianhaochen.springai.mongo.ChatMessageRepository;
import com.tianhaochen.springai.mongo.ChatRecordRepository;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.UUID;

@Component
@RequiredArgsConstructor
public class MongoChatMemory implements ChatMemory {
    private final ChatRecordRepository chatRecordRepository;
    private final ChatMessageRepository chatMessageRepository;

    @Override
    public void add(String conversationId, List<Message> messages) {

        List<ChatMessage> list = new ArrayList<>();

        messages.forEach(message -> list.add(new ChatMessage(UUID.randomUUID().toString(), conversationId, message)));
        chatMessageRepository.saveAll(list).subscribe();

    }

    @Override
    public List<Message> get(String conversationId, int lastN) {

        List<Message> messages = new ArrayList<>();
        chatMessageRepository.findAllByConversationId(conversationId).toStream()
                .forEach(chatMessage -> {
                    try {
                        switch (chatMessage.getMessageType()) {
                            case "USER" -> messages.add(new UserMessage(chatMessage.getTextContent()));
                            case "ASSISTANT" -> messages.add(new AssistantMessage(chatMessage.getTextContent()));
                            default ->
                                    throw new UnsupportedOperationException("Unsupported message type: " + chatMessage.getMessageType());
                        }
                    } catch (Exception e) {
                        // 处理异常
                    }
                });
        return messages != null ? messages.stream().skip(Math.max(0, messages.size() - lastN)).toList() : List.of();
    }

    @Override
    public void clear(String conversationId) {
        chatRecordRepository.deleteById(conversationId);
    }
}
